# ******************************************************************************
#  Copyright (c) 2021-2022. Kneron Inc. All rights reserved.                   *
# ******************************************************************************

# Legacy Module
from .KPValue import \
    GenericRawImageHeader, \
    GenericRawResult, \
    GenericRawBypassPreProcImageHeader, \
    GenericRawBypassPreProcResult

# Core Module
from ...KPValue import \
    DeviceGroup, \
    ModelNefDescriptor
from ...KPWrapper import KPWrapper
from ...KPUtils import _get_model_raw_out_size
from ...KPEnum import *
from ...KPException import _check_api_return_code
from typing import Union
import numpy as np
import ctypes


class KPInference:
    __KP_WRAPPER = KPWrapper()

    @staticmethod
    def generic_raw_inference_send(device_group: DeviceGroup,
                                   generic_raw_image_header: GenericRawImageHeader,
                                   image: Union[bytes, np.ndarray],
                                   image_format: ImageFormat) -> None:
        """
        Generic raw inference send.
        This is to perform a single image inference, it is non-blocking if device buffer queue is not full.
        When this is performed, user can issue kp.v1.inference.generic_raw_inference_receive() to get the result.
        In addition, to have better performance, users can issue multiple kp.v1.inference.generic_raw_inference_send() then
        start to receive results through kp.v1.inference.generic_raw_inference_receive().

        Parameters
        ----------
        device_group : kp.DeviceGroup
            Represents a set of devices handle.
        generic_raw_image_header : kp.v1.GenericRawImageHeader
            Needed parameters for performing inference including image width, height ..etc.
        image : bytes, numpy.ndarray
            The data bytes or numpy.ndarray contains the image.
        image_format : kp.ImageFormat
            Image format supported for inference.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.v1.inference.generic_raw_inference_receive : Generic raw inference receive.
        kp.v1.GenericRawImageHeader
        kp.ImageFormat
        """
        if isinstance(image, np.ndarray):
            generic_raw_image_header.width = image.shape[1]
            generic_raw_image_header.height = image.shape[0]
            image = image.tobytes()

        generic_raw_image_header.image_format = image_format

        image_buffer_c_array = (ctypes.c_ubyte * len(image)).from_buffer(bytearray(image))

        status = KPInference.__KP_WRAPPER.LIB.kp_generic_raw_inference_send(
            device_group.address,
            ctypes.byref(generic_raw_image_header._get_element_buffer()),
            image_buffer_c_array
        )

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def generic_raw_inference_receive(device_group: DeviceGroup,
                                      generic_raw_image_header: GenericRawImageHeader,
                                      model_nef_descriptor: ModelNefDescriptor) -> GenericRawResult:
        """
        Generic raw inference receive.
        When an image inference is done, this function can be used to get the results in RAW format.
        kp.inference.generic_inference_retrieve_float_node()/kp.inference.generic_inference_retrieve_fixed_node() to
        convert RAW format data to floating-point/fixed-point data.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            Represents a set of devices handle.
        generic_raw_image_header : kp.v1.GenericRawImageHeader
            Needed parameters for performing inference including image width, height ..etc.
        model_nef_descriptor : kp.ModelNefDescriptor
            ModelNefDescriptor object for describing the uploaded models.

        Returns
        -------
        generic_raw_result : kp.GenericRawResult
            GenericRawResult object contained the received RAW data results.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.v1.inference.generic_raw_inference_send : Generic raw inference send.
        kp.inference.generic_inference_retrieve_float_node : Retrieve single node output floating-point data buffer.
        kp.inference.generic_inference_retrieve_fixed_node : Retrieve single node output fixed-point data buffer.
        kp.v1.GenericRawImageHeader
        kp.ImageFormat
        kp.GenericRawResult

        Notes
        -----
        The data received is in Kneron RAW format, users need
        kp.inference.generic_inference_retrieve_float_node()/kp.inference.generic_inference_retrieve_fixed_node() to
        convert RAW format data to floating-point/fixed-point data.
        """
        raw_out_size = _get_model_raw_out_size(generic_inference_data_header=generic_raw_image_header,
                                               model_nef_descriptor=model_nef_descriptor)
        generic_raw_result = GenericRawResult(buffer_size=raw_out_size)

        status = KPInference.__KP_WRAPPER.LIB.kp_generic_raw_inference_receive(
            device_group.address,
            ctypes.byref(generic_raw_result.header._get_element_buffer()),
            generic_raw_result.raw_result._get_element_buffer()._LP_raw_out_buffer,
            raw_out_size
        )

        _check_api_return_code(result=generic_raw_result,
                               api_return_code=ApiReturnCode(status))

        return generic_raw_result

    @staticmethod
    def generic_raw_inference_bypass_pre_proc_send(device_group: DeviceGroup,
                                                   generic_raw_image_header: GenericRawBypassPreProcImageHeader,
                                                   image_buffer: bytes) -> None:
        """
        Generic raw inference bypass pre-processing send.
        This is to perform a single image inference, it is non-blocking if device buffer queue is not full. When this
        is performed, user can issue kp.inference.generic_raw_inference_bypass_pre_proc_receive() to get the result.
        In addition, to have better performance, users can issue multiple
        kp.inference.generic_raw_inference_bypass_pre_proc_send() then start to receive results through
        kp.inference.generic_raw_inference_bypass_pre_proc_receive().

        Parameters
        ----------
        device_group : kp.DeviceGroup
            Represents a set of devices handle.
        generic_raw_image_header : kp.GenericRawBypassPreProcImageHeader
            Needed parameters for performing bypass pre-processing inference including image buffer size, model ID ..etc.
        image_buffer : bytes
            The data bytes contains the image.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.inference.generic_raw_inference_bypass_pre_proc_receive : Generic raw inference bypass pre-processing receive.
        kp.GenericRawBypassPreProcImageHeader
        """
        image_buffer_c_array = (ctypes.c_ubyte * len(image_buffer)).from_buffer(bytearray(image_buffer))

        status = KPInference.__KP_WRAPPER.LIB.kp_generic_raw_inference_bypass_pre_proc_send(
            device_group.address,
            ctypes.byref(generic_raw_image_header._get_element_buffer()),
            image_buffer_c_array
        )

        _check_api_return_code(result=None,
                               api_return_code=ApiReturnCode(status))

    @staticmethod
    def generic_raw_inference_bypass_pre_proc_receive(device_group: DeviceGroup,
                                                      generic_raw_image_header: GenericRawBypassPreProcImageHeader,
                                                      model_nef_descriptor: ModelNefDescriptor) -> GenericRawBypassPreProcResult:
        """
        Generic raw inference bypass pre-processing receive.
        When an image inference is done, this function can be used to get the results in RAW format.
        Note that data received is in Kneron RAW format, users need
        kp.inference.generic_inference_retrieve_float_node()/kp.inference.generic_inference_retrieve_fixed_node() to
        convert RAW format data to floating-point/fixed-point data.

        Parameters
        ----------
        device_group : kp.DeviceGroup
            Represents a set of devices handle.
        generic_raw_image_header : kp.GenericRawBypassPreProcImageHeader
            Needed parameters for performing bypass pre-processing inference including image buffer size, model ID ... etc.
        model_nef_descriptor : kp.ModelNefDescriptor
            ModelNefDescriptor object for describing the uploaded models.

        Returns
        -------
        generic_raw_result : kp.GenericRawBypassPreProcResult
            GenericRawBypassPreProcResult object contained the received RAW data results.

        Raises
        ------
        kp.ApiKPException

        See Also
        --------
        kp.core.connect_devices : To connect multiple (including one) Kneron devices.
        kp.inference.generic_raw_inference_bypass_pre_proc_send : Generic raw inference bypass pre-processing send.
        kp.inference.generic_inference_retrieve_float_node : Retrieve single node output floating-point data buffer.
        kp.inference.generic_inference_retrieve_fixed_node : Retrieve single node output fixed-point data buffer.
        kp.GenericRawBypassPreProcImageHeader
        kp.ModelNefDescriptor
        kp.GenericRawBypassPreProcResult

        Notes
        -----
        The data received is in Kneron RAW format, users need
        kp.inference.generic_inference_retrieve_float_node()/kp.inference.generic_inference_retrieve_fixed_node() to
        convert RAW format data to floating-point/fixed-point data.
        """
        raw_out_size = _get_model_raw_out_size(generic_inference_data_header=generic_raw_image_header,
                                               model_nef_descriptor=model_nef_descriptor)
        generic_raw_bypass_pre_proc_result = GenericRawBypassPreProcResult(buffer_size=raw_out_size)

        status = KPInference.__KP_WRAPPER.LIB.kp_generic_raw_inference_bypass_pre_proc_receive(
            device_group.address,
            ctypes.byref(generic_raw_bypass_pre_proc_result.header._get_element_buffer()),
            generic_raw_bypass_pre_proc_result.raw_result._get_element_buffer()._LP_raw_out_buffer,
            raw_out_size
        )

        _check_api_return_code(result=generic_raw_bypass_pre_proc_result,
                               api_return_code=ApiReturnCode(status))

        return generic_raw_bypass_pre_proc_result

